import pandas as pd
import numpy as np
from statsmodels.tsa.statespace.sarimax import SARIMAX
from statsmodels.tsa.stattools import adfuller
import logging

logger = logging.getLogger(__name__)


class ARIMAModel:
    """ARIMA/SARIMA统计模型"""

    def __init__(self):
        self.model = None
        self.model_fit = None

    def test_stationarity(self, timeseries):
        """检验时间序列平稳性"""
        result = adfuller(timeseries, autolag='AIC')
        return result[1] < 0.05  # p-value < 0.05表示平稳

    def predict(self, df, forecast_hours=24, confidence_level=0.95):
        """
        使用SARIMA进行预测

        Args:
            df: DataFrame with columns ['timestamp', 'value']
            forecast_hours: 预测未来多少小时
            confidence_level: 置信区间水平

        Returns:
            预测结果列表
        """
        try:
            logger.info("开始SARIMA预测...")

            # 准备时间序列数据
            ts = pd.Series(
                df['value'].values,
                index=pd.to_datetime(df['timestamp'])
            )

            # 检验平稳性
            if not self.test_stationarity(ts):
                logger.warning("时间序列非平稳，将进行差分处理")

            # SARIMA模型参数
            # (p,d,q) - ARIMA参数
            # (P,D,Q,s) - 季节性参数，s=24表示24小时周期
            order = (2, 1, 2)  # ARIMA(2,1,2)
            seasonal_order = (1, 1, 1, 24)  # 24小时季节性

            # 训练SARIMA模型
            self.model = SARIMAX(
                ts,
                order=order,
                seasonal_order=seasonal_order,
                enforce_stationarity=False,
                enforce_invertibility=False
            )

            self.model_fit = self.model.fit(disp=False, maxiter=100)

            # 执行预测
            forecast_result = self.model_fit.get_forecast(steps=forecast_hours)
            predictions = forecast_result.predicted_mean
            conf_int = forecast_result.conf_int(alpha=1 - confidence_level)

            # 生成未来时间戳
            last_timestamp = pd.to_datetime(df['timestamp'].iloc[-1])
            future_timestamps = pd.date_range(
                start=last_timestamp + pd.Timedelta(hours=1),
                periods=forecast_hours,
                freq='H'
            )

            # 构建返回结果
            results = []
            for i, ts in enumerate(future_timestamps):
                pred_value = max(0, int(predictions.iloc[i]))
                results.append({
                    'timestamp': ts.strftime('%Y-%m-%d %H:%M:%S'),
                    'predicted_value': pred_value,
                    'confidence_lower': max(0, int(conf_int.iloc[i, 0])),
                    'confidence_upper': max(0, int(conf_int.iloc[i, 1]))
                })

            logger.info(f"SARIMA预测完成，生成{len(results)}个预测点")
            return results

        except Exception as e:
            logger.error(f"SARIMA预测失败: {str(e)}")
            raise